#!/usr/bin/env python3
import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os
import threading
from datetime import datetime

class ImageSaver:
    def __init__(self):
        self.bridge = CvBridge()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        
        # 配置参数
        self.save_dir = rospy.get_param("~save_dir", "/root/ros_ws/src/OrbbecSDK_ROS1/scripts/data")
        self.max_files = rospy.get_param("~max_files", 30)  # 最大保存文件数
        self.color_dir = os.path.join(self.save_dir, "Color3")
        self.depth_dir = os.path.join(self.save_dir, "Depth3")
        
        # 创建保存目录
        os.makedirs(self.color_dir, exist_ok=True)
        os.makedirs(self.depth_dir, exist_ok=True)
        
        # 清理旧文件
        self.cleanup_old_files()
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)
        rospy.loginfo("ImageSaver initialized, waiting for images...")
        
        # 启动文件清理定时器
        self.cleanup_timer = rospy.Timer(rospy.Duration(60), self.cleanup_old_files)

    def color_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            with self.lock:
                self.color_img = cv_image
                self.try_save_images()
        except Exception as e:
            rospy.logerr(f"Color error: {e}")

    def depth_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "passthrough")
            with self.lock:
                self.depth_img = cv_image
                self.try_save_images()
        except Exception as e:
            rospy.logerr(f"Depth error: {e}")

    def try_save_images(self):
        if self.color_img is not None and self.depth_img is not None:
            # 生成时间戳
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            color_path = os.path.join(self.color_dir, f"{timestamp}.png")
            depth_path = os.path.join(self.depth_dir, f"{timestamp}.png")
            
            # 保存图像
            cv2.imwrite(color_path, self.color_img)
            cv2.imwrite(depth_path, self.depth_img)
            
            rospy.loginfo_once(f"Saved first image pair: {timestamp}")
            rospy.loginfo_throttle(5,f"Saved image pair: {timestamp}")
            
            # 重置，避免重复保存同一对图像
            self.color_img = None
            self.depth_img = None
            
            # 检查文件数量
            self.cleanup_old_files()

    def cleanup_old_files(self, event=None):
        """清理旧文件，保持文件夹大小在限制范围内"""
        try:
            # 获取所有文件
            color_files = sorted(os.listdir(self.color_dir))
            depth_files = sorted(os.listdir(self.depth_dir))
            
            # 计算需要删除的文件数量
            total_files = len(color_files)
            if total_files > self.max_files:
                num_to_delete = total_files - self.max_files
                
                # 删除最旧的文件
                for i in range(num_to_delete):
                    os.remove(os.path.join(self.color_dir, color_files[i]))
                    os.remove(os.path.join(self.depth_dir, depth_files[i]))
                
                rospy.loginfo(f"Deleted {num_to_delete} old image pairs")
        except Exception as e:
            rospy.logerr(f"Cleanup error: {e}")

    def shutdown(self):
        """关闭时清理"""
        self.cleanup_timer.shutdown()
        rospy.loginfo("ImageSaver shutdown")

if __name__ == "__main__":
    rospy.init_node("image_saver")
    saver = ImageSaver()
    rospy.on_shutdown(saver.shutdown)
    rospy.spin()